#!/usr/bin/env python3
"""Remove the first row from a Parquet file, back up the original, and overwrite.

Usage:
    python scripts/remove_first_row_parquet.py /absolute/path/to/file.parquet

The script will create a backup under data/backup/ with a timestamp.
"""
from __future__ import annotations

import sys
import os
import shutil
from datetime import datetime
from pathlib import Path

try:
    import polars as pl
except Exception:
    pl = None

try:
    import pandas as pd
except Exception:
    pd = None


def backup_file(src: Path) -> Path:
    backup_dir = src.parent / "backup"
    backup_dir.mkdir(parents=True, exist_ok=True)
    ts = datetime.now().strftime("%Y%m%dT%H%M%S")
    dest = backup_dir / f"{src.name}.{ts}.bak"
    # copy the file instead of moving so the original remains for processing
    shutil.copy2(src, dest)
    return dest


def restore_backup(backup: Path, dest: Path) -> None:
    # copy backup back to destination (overwrite)
    shutil.copy2(backup, dest)


def remove_first_row_polars(path: Path) -> int:
    df = pl.read_parquet(path)
    if df.height == 0:
        return 0
    # drop first row by slicing
    df_new = df[1:]
    # overwrite parquet
    df_new.write_parquet(path)
    return df_new.height


def remove_first_row_pandas(path: Path) -> int:
    df = pd.read_parquet(path)
    if df.shape[0] == 0:
        return 0
    df_new = df.iloc[1:].reset_index(drop=True)
    df_new.to_parquet(path, index=False)
    return df_new.shape[0]


def main():
    if len(sys.argv) < 2:
        print("Usage: remove_first_row_parquet.py /path/to/file.parquet")
        sys.exit(2)

    target = Path(sys.argv[1])
    if not target.exists():
        print(f"File not found: {target}")
        sys.exit(2)

    # Make absolute
    target = target.resolve()

    # Backup original
    try:
        backup = backup_file(target)
        print(f"Backed up original to: {backup}")
    except Exception as e:
        print(f"Failed to backup file: {e}")
        sys.exit(1)

    try:
        if pl is not None:
            new_count = remove_first_row_polars(target)
            print(f"Wrote new parquet with {new_count} rows (Polars).")
        elif pd is not None:
            new_count = remove_first_row_pandas(target)
            print(f"Wrote new parquet with {new_count} rows (pandas).")
        else:
            print("Neither polars nor pandas is available to process parquet files.")
            restore_backup(backup, target)
            sys.exit(1)
    except Exception as e:
        print(f"Processing failed: {e}")
        # attempt restore
        try:
            restore_backup(backup, target)
            print("Original restored from backup.")
        except Exception as e2:
            print(f"Failed to restore backup: {e2}")
        sys.exit(1)


if __name__ == "__main__":
    main()
